//
//  erza_matrix_lib.c
//  KFilter
//
//  Created by Mita Yun on 6/20/13.
//  Copyright (c) 2013 com.mika. All rights reserved.
//

#include <stdio.h>
#include <stdlib.h>

#include "erza_matrix_lib.h"

#undef DEBUG
#define MIN(a, b) ((a < b) ? a : b)

void matrixInit(arm_matrix_instance_f32 *A, int rows, int cols)
{
	float32_t *d;

	d = (float32_t *) calloc(rows * cols, sizeof(float64_t));

	arm_mat_init_f32(A, rows, cols, d);
}

void matrixFree(arm_matrix_instance_f32 *A)
{
	if (A && A->pData)
		free(A->pData);
}

// Calculates the QR decomposition of the given matrix A Transposed (decomp's A', not A)
//      notes:  A matrix is modified
//      Adapted from Java code originaly written by Joni Salonen
//
// returns 1 for success, 0 for failure
int qrDecompositionT_f32(arm_matrix_instance_f32 *A, arm_matrix_instance_f32 *Q, arm_matrix_instance_f32 *R)
{
	int minor;
	int row, col;
	int m = A->numCols;
	int n = A->numRows;
	int min;

	// clear R
	arm_fill_f32(0, R->pData, R->numRows * R->numCols);

	min = MIN(m, n);

	/*
	 * The QR decomposition of a matrix A is calculated using Householder
	 * reflectors by repeating the following operations to each minor
	 * A(minor,minor) of A:
	 */
	for (minor = 0; minor < min; minor++)
	{
		float xNormSqr = 0.0f;
		float a;

		/*
		 * Let x be the first column of the minor, and a^2 = |x|^2.
		 * x will be in the positions A[minor][minor] through A[m][minor].
		 * The first column of the transformed minor will be (a,0,0,..)'
		 * The sign of a is chosen to be opposite to the sign of the first
		 * component of x. Let's find a:
		 */
		for (row = minor; row < m; row++)
			xNormSqr += A->pData[minor * m + row] * A->pData[minor * m + row];

		a = sqrtf(xNormSqr);
		if (A->pData[minor * m + minor] > 0.0f)
			a = -a;

		if (a != 0.0f)
		{
			R->pData[minor * R->numCols + minor] = a;

			/*
			 * Calculate the normalized reflection vector v and transform
			 * the first column. We know the norm of v beforehand: v = x-ae
			 * so |v|^2 = <x-ae,x-ae> = <x,x>-2a<x,e>+a^2<e,e> =
			 * a^2+a^2-2a<x,e> = 2a*(a - <x,e>).
			 * Here <x, e> is now A[minor][minor].
			 * v = x-ae is stored in the column at A:
			 */
			A->pData[minor * m + minor] -= a; // now |v|^2 = -2a*(A[minor][minor])

			/*
			 * Transform the rest of the columns of the minor:
			 * They will be transformed by the matrix H = I-2vv'/|v|^2.
			 * If x is a column vector of the minor, then
			 * Hx = (I-2vv'/|v|^2)x = x-2vv'x/|v|^2 = x - 2<x,v>/|v|^2 v.
			 * Therefore the transformation is easily calculated by
			 * subtracting the column vector (2<x,v>/|v|^2)v from x.
			 *
			 * Let 2<x,v>/|v|^2 = alpha. From above we have
			 * |v|^2 = -2a*(A[minor][minor]), so
			 * alpha = -<x,v>/(a*A[minor][minor])
			 */
			for (col = minor + 1; col < n; col++)
			{
				float alpha = 0.0f;

				for (row = minor; row < m; row++)
					alpha -= A->pData[col * m + row] * A->pData[minor * m + row];

				alpha /= a * A->pData[minor * m + minor];

				// Subtract the column vector alpha*v from x.
				for (row = minor; row < m; row++)
					A->pData[col * m + row] -= alpha * A->pData[minor * m + row];
			}
		}
		// rank deficient
		else
			return 0;
	}

	// Form the matrix R of the QR-decomposition.
	//      R is supposed to be m x n, but only calculate n x n
	// copy the upper triangle of A
	for (row = min - 1; row >= 0; row--)
		for (col = row + 1; col < n; col++)
			R->pData[row * R->numCols + col] = A->pData[col * m + row];

	// Form the matrix Q of the QR-decomposition.
	//      Q is supposed to be m x m

	// only compute Q if requested
	if (Q)
	{
		arm_fill_f32(0, Q->pData, Q->numRows * Q->numCols);

		/*
		 * Q = Q1 Q2 ... Q_m, so Q is formed by first constructing Q_m and then
		 * applying the Householder transformations Q_(m-1),Q_(m-2),...,Q1 in
		 * succession to the result
		 */
		for (minor = m - 1; minor >= min; minor--)
			Q->pData[minor * m + minor] = 1.0f;

		for (minor = min - 1; minor >= 0; minor--)
		{
			Q->pData[minor * m + minor] = 1.0f;

			if (A->pData[minor * m + minor] != 0.0f)
			{
				for (col = minor; col < m; col++)
				{
					float alpha = 0.0f;

					for (row = minor; row < m; row++)
						alpha -= Q->pData[row * m + col] * A->pData[minor * m + row];

					alpha /= R->pData[minor * R->numCols + minor] * A->pData[minor * m + minor];

					for (row = minor; row < m; row++)
						Q->pData[row * m + col] -= alpha * A->pData[minor * m + row];
				}
			}
		}
	}

	return 1;
}

// Solves m sets of n equations B * X = A using QR decomposition and backsubstitution
void matrixDiv_f32(arm_matrix_instance_f32 *X, arm_matrix_instance_f32 *A, arm_matrix_instance_f32 *B, arm_matrix_instance_f32 *Q, arm_matrix_instance_f32 *R, arm_matrix_instance_f32 *AQ) {
	int i, j, k;
	int m, n;

	// this is messy (going into a class's private data structure),
	// but it is better than malloc/free
	Q->numRows = B->numRows;
	Q->numCols = B->numRows;
	R->numRows = B->numRows;
	R->numCols = B->numCols;
	AQ->numRows = A->numRows;
	AQ->numCols = B->numRows;

	m = A->numRows;
	n = B->numCols;

	qrDecompositionT_f32(B, Q, R);
	arm_mat_mult_f32(A, Q, AQ);

	// solve for X by backsubstitution
	for (i = 0; i < m; i++) {
		for (j = n-1; j >= 0; j--) {
			for (k = j+1; k < n; k++)
				AQ->pData[i*n + j] -= R->pData[j*n + k] * X->pData[i*n + k];
			X->pData[i*n + j] = AQ->pData[i*n + j] / R->pData[j*n + j];
		}
	}
}


/*QR decomp using gramt-schmidt*/
arm_status qr(arm_matrix_instance_f32 *A, arm_matrix_instance_f32 *Q, arm_matrix_instance_f32 *R)
{
	int row, col, i, j;
	int m = A->numRows;
	int n = A->numCols;

	arm_matrix_instance_f32 A_t;
	matrixInit(&A_t, n, m);
	arm_mat_trans_f32(A, &A_t);

#ifdef DEBUG
	printf("A_t: \n");
	matrixPrint(&A_t);
#endif

	if (R->numRows != A->numRows || R->numCols != A->numCols || Q->numCols != A->numRows || Q->numRows != A->numRows || A->numRows < A->numCols)
	{
		printf("argument error");
		return ARM_MATH_ARGUMENT_ERROR;
	}

	arm_matrix_instance_f32 Q_t;
	matrixInit(&Q_t, m, m);

	arm_matrix_instance_f32 Q_hat_t;
	matrixInit(&Q_hat_t, m, m);

	// Q_hat_t[0] = A_t[0];
	float32_t sum = 0;
	for (j = 0; j < Q_hat_t.numCols; j++)
	{
		float32_t data = matrixGet(&A_t, 0, j);
		matrixSet(data, &Q_hat_t, 0, j);
		sum += data * data;
	}

	// Q_t[0] = multiply(1 / abs(Q_hat_t[0]), Q_hat_t[0])
	for (j = 0; j < Q_t.numCols; j++)
	{
		float32_t data = matrixGet(&Q_hat_t, 0, j);
		data /= sqrtf(sum);
		matrixSet(data, &Q_t, 0, j);
	}

#ifdef DEBUG
	printf("initial Q_t: \n");
	matrixPrint(&Q_t);
#endif

	/*
	 R[0][0] = abs(A_t[0]);
	 for (int col = 1; col < n; col++) {
	 R[0][col] = dot(Q_t[0], A_t[col]);
	 }
	 */
	matrixSet(vectorMag(A_t.pData, 0, A_t.numCols), R, 0, 0);
	for (j = 1; j < R->numCols; j++)
	{
		sum = 0;
		for (int k = 0; k < A_t.numCols; k++)
		{
			sum += matrixGet(&Q_t, 0, k) * matrixGet(&A_t, j, k);
		}
		matrixSet(sum, R, 0, j);
	}
#ifdef DEBUG
	printf("initial R: \n");
	matrixPrint(R);
#endif

	float32_t *temp_n = (float32_t *) calloc(A_t.numCols, sizeof(float32_t));
	for (row = 1; row < m; row++)
	{
		if (row >= n)
		{
			break;
		}

		//temp_n = A_t[row];
		vectorCopy(A_t.pData, temp_n, row * A_t.numCols, (row + 1) * A_t.numCols);

		/*
		 for (int hat_inner = row - 1; hat_inner >= 0; hat_inner--) {
		 temp_n = subtract(temp_n,
		 multiply(R[hat_inner][row], Q_t[hat_inner]));
		 }
		 */
		for (int hat_inner = row - 1; hat_inner >= 0; hat_inner--)
		{
			for (i = 0; i < A_t.numCols; i++)
			{
				temp_n[i] -= matrixGet(R, hat_inner, row) * matrixGet(&Q_t, hat_inner, i);
			}
		}
#ifdef DEBUG
		vectorPrint(temp_n, 0, A_t.numCols);
#endif

		/*
		 Q_hat_t[row] = temp_n;
		 Q_t[row] = multiply(1 / abs(Q_hat_t[row]), Q_hat_t[row]);
		 */
		sum = 0;
		for (j = 0; j < Q_hat_t.numCols; j++)
		{
			matrixSet(temp_n[j], &Q_hat_t, row, j);
			sum += matrixGet(&Q_hat_t, row, j) * matrixGet(&Q_hat_t, row, j);
		}
#ifdef DEBUG
		printf("Q_hat_t\n");
#endif
		matrixPrint(&Q_hat_t);
		for (j = 0; j < Q_t.numCols; j++)
		{
			matrixSet(matrixGet(&Q_hat_t, row, j) / sqrtf(sum), &Q_t, row, j);
		}
#ifdef DEBUG
		printf("Q_t\n");
#endif
		matrixPrint(&Q_t);

		for (col = 1; col < n; col++)
		{
			if (row == col)
			{
				// R[row][col] = abs(Q_hat_t[col]);
				sum = 0;
				for (int j = 0; j < Q_hat_t.numCols; j++)
				{
					sum += matrixGet(&Q_hat_t, col, j) * matrixGet(&Q_hat_t, col, j);
				}
				matrixSet(sqrtf(sum), R, row, col);
			}
			else
			{
				// R[row][col] = dot(Q_t[row], A_t[col]);
				sum = 0;
				for (int j = 0; j < Q_t.numCols; j++)
				{
					sum += matrixGet(&Q_t, row, j) * matrixGet(&A_t, col, j);
				}
				matrixSet(sum, R, row, col);
			}
		}

	}

	arm_mat_trans_f32(&Q_t, Q);

	// matrixPrint(R);
	//    matrixPrint(Q);
	return 1;
}

void vectorPrint(float32_t *x, int start, int end)
{
	printf("\n [");
	for (int i = start; i < end; i++)
	{
		printf("%f, ", x[i]);
	}
	printf("] \n");
}

void vectorCopy(float32_t *src, float32_t *dest, int start, int end)
{
	int s = 0;
	for (int i = start; i < end; i++)
	{
		dest[s++] = src[i];
	}
}

float32_t vectorMag(float32_t *x, int start, int end)
{
	float32_t sum = 0;
	for (int i = start; i < end; i++)
	{
		sum += x[i] * x[i];
	}
	return sqrtf(sum);
}

void vectorNorm(float32_t *src, float32_t *dest, int start, int end)
{
	float32_t mag = vectorMag(src, start, end);
	for (int i = start; i < end; i++)
	{
		dest[i] = src[i] / mag;
	}
}

float32_t matrixGet(arm_matrix_instance_f32 *A, int rowIndex, int colIndex)
{
	if (rowIndex >= A->numRows || colIndex >= A->numCols)
	{
		printf("argument error");
		return ARM_MATH_ARGUMENT_ERROR;
	}
	return A->pData[rowIndex * A->numCols + colIndex];
}

void matrixSet(float32_t data, arm_matrix_instance_f32 *A, int rowIndex, int colIndex)
{
	if (rowIndex >= A->numRows || colIndex >= A->numCols)
	{
		// printf("argument error");
		return;
	}
	A->pData[rowIndex * A->numCols + colIndex] = data;
}
void mprintf(char * name, arm_matrix_instance_f32 *A)
{
	int m = A->numRows;
	int n = A->numCols;
	printf("--------%s (%d,%d)--------\n", name, m, n);

	for (int i = 0; i < m; i++)
	{
		printf("[");
		for (int j = 0; j < n; j++)
		{
			float f = A->pData[i * n + j];
			printf("%12g, ", f);
		}
		printf("]\n");
	}
	printf("--------\n");

}

void matrixPrint(arm_matrix_instance_f32 *A)
{
	int m = A->numRows;
	int n = A->numCols;
	printf("start printing %d by %d matrix: \n", m, n);

	for (int i = 0; i < m; i++)
	{
		printf("[");
		for (int j = 0; j < n; j++)
		{
			float f = A->pData[i * n + j];
			printf("%12f, ", f);
		}
		printf("]\n");
	}
	printf("end printing matrix. \n");
}


/*Other matrix helpers*/

void quatNorm(float *inQuat, float *outQuat) {
    float norm = sqrtf(inQuat[0]*inQuat[0] + inQuat[1]*inQuat[1] + inQuat[2]*inQuat[2] + inQuat[3]*inQuat[3]);

    outQuat[0] = inQuat[0] / norm;
    outQuat[1] = inQuat[1] / norm;
    outQuat[2] = inQuat[2] / norm;
    outQuat[3] = inQuat[3] / norm;
}


void quatToEuler(float *quat, float *outYaw, float *outPitch, float *outRoll) {

	float q0 = quat[0];
	float q1 = quat[1];
	float q2 = quat[2];
	float q3 = quat[3];

	*outYaw = atan2(2*(q0*q3+q1*q2), 1-2*(q2*q2+q3*q3));
	*outPitch = asin(2*(q0*q2-q3*q1));
	*outRoll = atan2(2*(q0*q1+q2*q3), 1-2*(q1*q1+q2*q2));
}
